# RUN: tf-mlir-translate -graphdef-to-mlir -tf-enable-shape-inference-on-import=false %s -tf-input-arrays=a,b -tf-input-data-types=DT_FLOAT,DT_FLOAT -tf-input-shapes=':' -tf-output-arrays=StatefulIf,StatelessIf -o - -mlir-print-debuginfo -mlir-print-local-scope | FileCheck %s

# Verify that TensorFlow If and StatelessIf ops are mapped to the
# composite If op in MLIR with is_stateless attribute set accordingly to
# distinguish between them.

# CHECK-DAG: "tf.If"{{.*}} is_stateless = false{{.*}} loc(fused["If:", "StatefulIf"])
# CHECK-DAG: "tf.If"{{.*}} is_stateless = true{{.*}} loc(fused["StatelessIf:", "StatelessIf"])

node {
  name: "tf.Less"
  op: "Less"
  input: "a"
  input: "b"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  experimental_debug_info {
  }
}
node {
  name: "StatefulIf"
  op: "If"
  input: "tf.Less"
  input: "a"
  input: "b"
  attr {
    key: "Tcond"
    value {
      type: DT_BOOL
    }
  }
  attr {
    key: "Tin"
    value {
      list {
        type: DT_FLOAT
        type: DT_FLOAT
      }
    }
  }
  attr {
    key: "Tout"
    value {
      list {
        type: DT_FLOAT
      }
    }
  }
  attr {
    key: "else_branch"
    value {
      func {
        name: "cond_false"
      }
    }
  }
  attr {
    key: "then_branch"
    value {
      func {
        name: "cond_true"
      }
    }
  }
  experimental_debug_info {
  }
}
node {
  name: "StatelessIf"
  op: "StatelessIf"
  input: "tf.Less"
  input: "a"
  input: "b"
  attr {
    key: "Tcond"
    value {
      type: DT_BOOL
    }
  }
  attr {
    key: "Tin"
    value {
      list {
        type: DT_FLOAT
        type: DT_FLOAT
      }
    }
  }
  attr {
    key: "Tout"
    value {
      list {
        type: DT_FLOAT
      }
    }
  }
  attr {
    key: "else_branch"
    value {
      func {
        name: "cond_false"
      }
    }
  }
  attr {
    key: "then_branch"
    value {
      func {
        name: "cond_true"
      }
    }
  }
  experimental_debug_info {
  }
}
node {
  name: "main"
  op: "_Retval"
  input: "StatefulIf"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "index"
    value {
      i: 0
    }
  }
}
node {
  name: "main1"
  op: "_Retval"
  input: "StatelessIf"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "index"
    value {
      i: 1
    }
  }
}
node {
  name: "a"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  experimental_debug_info {
  }
}
node {
  name: "b"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  experimental_debug_info {
  }
}
library {
  function {
    signature {
      name: "cond_true"
      input_arg {
        name: "cond_true"
        type: DT_FLOAT
      }
      input_arg {
        name: "cond_true1"
        type: DT_FLOAT
      }
      output_arg {
        name: "cond_true2"
        type: DT_FLOAT
      }
    }
    node_def {
      name: "tf.Add"
      op: "Add"
      input: "cond_true"
      input: "cond_true1"
      attr {
        key: "T"
        value {
          type: DT_FLOAT
        }
      }
      experimental_debug_info {
        original_node_names: "tf.Add"
      }
    }
    ret {
      key: "cond_true2"
      value: "tf.Add:z:0"
    }
  }
  function {
    signature {
      name: "cond_false"
      input_arg {
        name: "cond_false"
        type: DT_FLOAT
      }
      input_arg {
        name: "cond_false1"
        type: DT_FLOAT
      }
      output_arg {
        name: "cond_false2"
        type: DT_FLOAT
      }
    }
    node_def {
      name: "tf.Mul"
      op: "Mul"
      input: "cond_false"
      input: "cond_false1"
      attr {
        key: "T"
        value {
          type: DT_FLOAT
        }
      }
      experimental_debug_info {
        original_node_names: "tf.Mul"
      }
    }
    ret {
      key: "cond_false2"
      value: "tf.Mul:z:0"
    }
  }
}
versions {
  producer: 115
  min_consumer: 12
}

